import pymysql
import json
import os
import re
from typing import Dict, List, Any, Optional
from datetime import datetime
from pymysql.cursors import DictCursor
import yaml

# 读取配置文件
def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

# 获取CVE数据文件夹路径
config = load_config()
CVE_DATA_FOLDER = config['cve']['data_folder']

class MySQLDefectDetector:
    def __init__(self, host: str, user: str, password: str, port: int = 3306, database: str = None):
        self.host = host
        self.user = user
        self.password = password
        self.port = port
        self.database = database
        self.major_version = None
        self.full_version = None
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Accept": "application/json",
        }
        self.conn = None
        self.cursor = None

    def _create_temp_connection(self) -> bool:
        try:
            self.conn = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                port=self.port,
                cursorclass=DictCursor,
                charset='utf8mb4',
                connect_timeout=5
            )
            self.cursor = self.conn.cursor()
            self.cursor.execute("SELECT VERSION() as version")
            version_result = self.cursor.fetchone()
            if version_result and version_result.get('version'):
                self.full_version = version_result['version']
                version_parts = self.full_version.split('.')
                if len(version_parts) >= 1:
                    self.major_version = version_parts[0]
                if self.database:
                    self.use_database(self.database)
            return True
        except pymysql.Error as e:
            print(f"临时连接失败: {str(e)}")
            return False

    def _close_temp_connection(self) -> None:
        try:
            if self.cursor:
                self.cursor.close()
                self.cursor = None
            if self.conn and self.conn.open:
                self.conn.close()
                self.conn = None
        except Exception as e:
            print(f"临时连接关闭失败: {e}")

    def use_database(self, database: str) -> None:
        if not self.cursor:
            print("无法切换数据库:未建立连接")
            return
        try:
            self.cursor.execute(f"USE `{database}`")
            print(f"已切换到数据库: {database}")
        except pymysql.Error as e:
            print(f"切换数据库失败: {e}")

    def check_version_vulnerabilities(self) -> Dict[str, Any]:
        if not self._create_temp_connection():
            return {"error": "无法创建连接获取版本信息", "version": "unknown", "vulnerabilities": []}

        try:
            version = self.full_version or "unknown"
            print(f"检测到 MySQL版本: {version}")

            # 加载本地CVE数据
            local_cves = self._load_local_cves()

            if not local_cves:
                print("未发现任何本地潜在漏洞")
                return {'version': version, 'vulnerabilities': []}

            confirmed_vulns = []

            for cve in local_cves:
                if not isinstance(cve, dict) or 'cve_id' not in cve:
                    print("跳过无效的 CVE条目")
                    continue

                print(f"\n=====测试 CVE: {cve['cve_id']}====")

                if self._test_single_cve(cve, version):
                    confirmed_vulns.append(cve)
                    print(f"[!] 确认存在漏洞: {cve['cve_id']}")
                else:
                    print(f"[√] 未发现漏洞: {cve['cve_id']}")

            return {
                'version': version,
                'vulnerabilities': confirmed_vulns
            }

        finally:
            self._close_temp_connection()

    def _load_local_cves(self) -> List[Dict[str, Any]]:
        if not os.path.exists(CVE_DATA_FOLDER):
            print(f"本地CVE文件夹 {CVE_DATA_FOLDER} 不存在")
            return []
        
        if not os.path.isdir(CVE_DATA_FOLDER):
            print(f"{CVE_DATA_FOLDER} 不是文件夹")
            return []
        
        all_cves = []
        
        # 遍历文件夹中的所有文件
        for filename in os.listdir(CVE_DATA_FOLDER):
            filepath = os.path.join(CVE_DATA_FOLDER, filename)
            
            # 只处理JSON文件
            if not filename.endswith('.json'):
                continue
                
            if not os.path.isfile(filepath):
                continue
                
            try:
                with open(filepath, 'r', encoding='utf-8') as f:
                    content = f.read()
                    if not content.strip():
                        print(f"文件 {filename} 为空，跳过")
                        continue
                    
                    file_data = json.loads(content)
                    if not isinstance(file_data, dict):
                        print(f"文件 {filename} 格式不正确，顶层不是对象，跳过")
                        continue
                    
                    # NVD格式的CVE数据在vulnerabilities字段中
                    if "vulnerabilities" in file_data and isinstance(file_data["vulnerabilities"], list):
                        vulnerabilities = file_data["vulnerabilities"]
                        for vuln in vulnerabilities:
                            if "cve" in vuln and "id" in vuln["cve"]:
                                cve_id = vuln["cve"]["id"]
                                # 从vulnerability中提取更多有用信息
                                description = ""
                                if "descriptions" in vuln and isinstance(vuln["descriptions"], list):
                                    for desc in vuln["descriptions"]:
                                        if desc.get("lang") == "en":
                                            description = desc.get("value", "")
                                            break
                                
                                # 构造简化版的CVE字典，包含我们需要的字段
                                simplified_cve = {
                                    "cve_id": cve_id,
                                    "description": description,
                                    # 可以添加更多需要的字段
                                    "published": vuln.get("published", ""),
                                    "lastModified": vuln.get("lastModified", ""),
                                    "vulnStatus": vuln.get("vulnStatus", "")
                                }
                                all_cves.append(simplified_cve)
                                print(f"从文件 {filename} 加载了 CVE: {cve_id}")
                    else:
                        print(f"文件 {filename} 缺少 'vulnerabilities' 字段或格式不正确，跳过")
                    
            except json.JSONDecodeError as e:
                print(f"文件 {filename} 不是有效的JSON格式: {e}，跳过")
            except Exception as e:
                print(f"读取文件 {filename} 时发生错误: {e}，跳过")
        
        print(f"总共加载了 {len(all_cves)} 个本地CVE条目")
        return all_cves

    def _test_single_cve(self, cve: Dict[str, Any], mysql_version: str) -> bool:
        if not isinstance(cve, dict):
            print("无效的CVE数据")
            return False

        if not self._is_cve_affected_by_version(cve, mysql_version):
            return False

        cve_id = cve['cve_id']
        if self.major_version == '5':
            fixed_cves = [
                "CVE-2016-6662", "CVE-2021-2154", "CVE-2018-3251",
                "CVE-2019-2737", "CVE-2020-14765", "CVE-2022-21363"
            ]
            if cve_id in fixed_cves:
                print(f"CVE {cve_id} 在 MySQL 5.7.44 中已修复，跳过测试")
                return False

        if not self._create_temp_connection():
            print("无法创建测试连接，跳过验证")
            return False

        try:
            print(f"执行漏洞验证逻辑 for {cve['cve_id']}")
            if cve_id == "CVE-2016-6662":
                return self._test_cve_2016_6662()
            elif cve_id == "CVE-2021-2154":
                return self._test_cve_2021_2154()
            elif cve_id == "CVE-2018-2755":
                return self._test_cve_2018_2755()
            elif cve_id == "CVE-2019-2232":
                return self._test_cve_2019_2232()
            else:
                # 对于未实现的CVE，可以根据需要添加测试逻辑
                print(f"CVE {cve_id} 的测试逻辑尚未实现，默认返回True")
                return True
        finally:
            self._close_temp_connection()

    def _test_cve_2016_6662(self) -> bool:
        try:
            self.cursor.execute("""
                CREATE FUNCTION mysql.user()
                RETURNS STRING
                SONAME 'ha_mysql.so'
            """)
            return True
        except pymysql.Error as e:
            if "ha_mysql.so: cannot open shared object file" in str(e):
                return False
            return True
        except Exception:
            return False

    def _test_cve_2021_2154(self) -> bool:
        try:
            payload = "A" * 1024
            self.cursor.execute(f'SELECT {payload}')
            return False  # 如果执行成功则说明可能不受影响
        except pymysql.OperationalError as e:
            if "Lost connection to MySQL server during query" in str(e):
                return True  # 连接丢失可能表示漏洞存在
            return False
        except Exception:
            return False

    def _test_cve_2018_2755(self) -> bool:
        try:
            self.cursor.execute("""
                SELECT @@GLOBAL.innodb_log_buffer_size;
                SET GLOBAL innodb_log_buffer_size=1;
            """)
            return True
        except pymysql.Error as e:
            if "Variable 'innodb_log_buffer_size' is a read-only variable" in str(e):
                return False
            return True
        except Exception:
            return False

    def _test_cve_2019_2232(self) -> bool:
        try:
            self.cursor.execute("""
                CREATE TABLE test_cve(a INT);
                INSERT INTO test_cve VALUES(1);
                SELECT /*!40001 SQL_NO_CACHE*/* FROM test_cve
                WHERE a=(SELECT a FROM test_cve GROUP BY a HAVING COUNT(*) = 1);
                DROP TABLE test_cve;
            """)
            return False
        except pymysql.OperationalError as e:
            if "Got error 139 from storage engine" in str(e):
                return True
            return False
        except Exception:
            return False

    def _is_cve_affected_by_version(self, cve_item: dict, mysql_version: str) -> bool:
        if not isinstance(cve_item, dict) or not mysql_version:
            return False
        try:
            version_match = re.match(r'^(\d+)\.(\d+)\.(\d+)', mysql_version)
            if not version_match:
                print(f"无法解析MySQL版本: {mysql_version}")
                return False
            current_parts = list(map(int, version_match.groups()))
        except (ValueError, TypeError):
            print(f"无法解析 MySQL版本: {mysql_version}")
            return False

        cve_config = cve_item.get("cve", {}).get("configurations", {})
        if not isinstance(cve_config, dict):
            return False

        for node in cve_config.get("nodes", []):
            for cpe_match in node.get("cpeMatch", []):
                criteria = cpe_match.get("criteria", "")
                if "mysql" not in criteria.lower():
                    continue
                version_match = re.search(r"mysql:([\d.]+)", criteria)
                if not version_match:
                    continue
                try:
                    cpe_version = version_match.group(1)
                    cpe_version_match = re.match(r'^(\d+)\.(\d+)\.(\d+)', cpe_version)
                    if not cpe_version_match:
                        continue
                    cpe_parts = list(map(int, cpe_version_match.groups()))
                except (ValueError, TypeError):
                    continue

                version_start = cpe_match.get('versionStartExcluding') or cpe_match.get('versionStartIncluding')
                version_end = cpe_match.get('versionEndExcluding') or cpe_match.get('versionEndIncluding')

                if version_start:
                    try:
                        start_match = re.match(r'^(\d+)\.(\d+)\.(\d+)', version_start)
                        if start_match:
                            start_parts = list(map(int, start_match.groups()))
                            if current_parts < start_parts:
                                continue
                    except (ValueError, TypeError):
                        continue
                if version_end:
                    try:
                        end_match = re.match(r'^(\d+)\.(\d+)\.(\d+)', version_end)
                        if end_match:
                            end_parts = list(map(int, end_match.groups()))
                            if current_parts > end_parts:
                                continue
                    except (ValueError, TypeError):
                        continue

                if self.major_version == '5' and current_parts[1] == 7:
                    if current_parts[0] == cpe_parts[0] and current_parts[1] == cpe_parts[1]:
                        return True
                    else:
                        return False
                return True
        return False

    def analyze_slow_queries(self) -> Dict[str, Any]:
        # 实现慢查询分析方法
        # 这里只是一个示例实现，返回一个示例字典
        try:
            # 查询慢查询日志或使用SHOW VARIABLES和SHOW STATUS获取相关信息
            # 这里只是一个示例，返回一个空字典
            return {
                'threshold': 1.0,  # 示例阈值
                'slow_queries': []  # 示例慢查询列表
            }
        except Exception as e:
            print(f"分析慢查询时出错: {e}")
            return {}

    def check_config_issues(self) -> List[str]:
        # 您的配置问题检查代码保持不变
        # 这里只是一个示例实现，返回空列表
        return []

    def analyze_table_design(self, table_name: str) -> List[str]:
        # 您的表设计分析代码保持不变
        # 这里只是一个示例实现，返回空列表
        return []

    def get_all_tables(self) -> List[str]:
        # 实现获取所有表的方法
        # 这里只是一个示例实现，返回空列表
        try:
            if not self.cursor:
                print("没有可用的数据库连接")
                return []
            self.cursor.execute("SHOW TABLES")
            tables = [row[f"Tables_in_{self.database}"] for row in self.cursor.fetchall()]
            return tables
        except Exception as e:
            print(f"获取所有表时出错: {e}")
            return []

    @staticmethod
    def generate_html_report(version_info: Dict[str, Any], slow_queries: Dict[str, Any], config_issues: List[str], table_issues: List[str] = None) -> None:
        # 您的HTML报告生成代码保持不变
        # 这里只是一个示例实现
        print("生成HTML报告...")

    def close(self) -> None:
        self._close_temp_connection()
        print("数据库连接已关闭")

def main() -> None:
    print("\n===MySQL数据库健康检测工具==")
    config = {
        'host': config['database']['host'],
        'user': config['database']['user'],
        'password': config['database']['password'],
        'port': config['database']['port'],
        'database': config['database']['database']
    }
    detector = MySQLDefectDetector(**config)
    try:
        print("\n[1/4]检测版本漏洞...")
        version_info = detector.check_version_vulnerabilities()
        if not isinstance(version_info, dict):
            print("版本检测返回无效数据")
            return
        print(f"版本: {version_info.get('version', '未知')}")
        print(f"确认漏洞数: {len(version_info.get('vulnerabilities', []))}")

        print("\n[2/4]分析慢查询...")
        slow_queries = detector.analyze_slow_queries()
        if isinstance(slow_queries, dict):
            print(f"阈值: {slow_queries.get('threshold', '未知')}秒")
            # 使用双引号包裹f-string，内部使用单引号
            print(f"慢查询数: {len(slow_queries.get('slow_queries', []))}")
        else:
            print("慢查询分析返回无效数据")

        print("\n[3/4]检查配置问题...")
        config_issues = detector.check_config_issues()
        if isinstance(config_issues, list):
            print(f"配置问题数: {len(config_issues)}")
        else:
            print("配置检查返回无效数据")

        print("\n[4/4]分析表结构...")
        all_table_issues = 0
        tables = detector.get_all_tables()
        if tables:
            print(f"发现{len(tables)}个表,开始分析...")
            for table in tables:
                issues = detector.analyze_table_design(table)
                all_table_issues.extend(issues)
        else:
            print("未发现任何表或无法获取表列表")

        print("\n生成报告...")
        MySQLDefectDetector.generate_html_report(version_info, slow_queries, config_issues, all_table_issues)

    except Exception as e:
        print(f"\n检测过程中发生错误: {e}")
    finally:
        detector.close()
        print("\n===检测完成==")

if __name__ == "__main__":
    main()
    